#!/usr/bin/env python3
"""
数据预处理运行脚本

在PyCharm中运行此脚本来执行数据预处理
"""

import sys
from pathlib import Path

# 添加项目路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root / 'src'))
sys.path.insert(0, str(project_root))

from scripts.preprocessing.preprocess_echo import EchoDataPreprocessor
from utils.logger import get_logger


def main():
    """主函数"""
    logger = get_logger("RunPreprocessing")
    
    print("=" * 60)
    print("SAM-Echo-Video 数据预处理")
    print("=" * 60)
    
    # 数据集路径
    datasets = {
        'BraTS2023': r"D:\Data\BraTS2023",
        'EchoNet-Dynamic': r"D:\Data\EchoNet-Dynamic", 
        'CardiacNet': r"D:\Data\CardiacNet"
    }
    
    # 检查数据集是否存在
    available_datasets = {}
    for name, path in datasets.items():
        if Path(path).exists():
            available_datasets[name] = path
            print(f"✓ {name}: {path}")
        else:
            print(f"✗ {name}: {path} (不存在)")
    
    if not available_datasets:
        print("\n错误: 没有找到可用的数据集!")
        print("请确保数据集路径正确，或手动修改脚本中的路径。")
        return False
    
    print(f"\n找到 {len(available_datasets)} 个可用数据集")
    
    # 创建预处理器
    try:
        preprocessor = EchoDataPreprocessor(
            config_path="configs/training.yaml",
            output_dir="data/processed/segmentation"
        )
        print("✓ 预处理器初始化成功")
    except Exception as e:
        print(f"✗ 预处理器初始化失败: {e}")
        return False
    
    # 运行预处理
    try:
        print("\n开始数据预处理...")
        print("使用快速模式：只处理部分数据用于测试")
        results = preprocessor.run_preprocessing(
            brats_dir=available_datasets.get('BraTS2023'),
            echonet_dir=available_datasets.get('EchoNet-Dynamic'),
            cardiacnet_dir=available_datasets.get('CardiacNet'),
            fast_mode=True  # 启用快速模式
        )
        
        print("\n" + "=" * 60)
        print("数据预处理完成!")
        print("=" * 60)
        
        # 打印结果摘要
        total_samples = 0
        for dataset, result in results.items():
            if isinstance(result, dict) and 'valid_samples' in result:
                valid = result['valid_samples']
                total = result['total_samples']
                print(f"{dataset}: {valid}/{total} 个样本")
                total_samples += valid
        
        print(f"\n总计: {total_samples} 个有效样本")
        
        # 打印数据划分
        if 'splits' in results:
            print("\n数据划分:")
            for split_name, samples in results['splits'].items():
                print(f"  {split_name}: {len(samples)} 个样本")
        
        print(f"\n处理结果保存到: data/processed/segmentation/")
        print("下一步: 运行模型训练")
        
        return True
        
    except Exception as e:
        print(f"✗ 数据预处理失败: {e}")
        logger.error(f"数据预处理失败: {e}")
        return False


if __name__ == "__main__":
    success = main()
    if not success:
        print("\n数据预处理失败，请检查错误信息并重试。")
        sys.exit(1)
    else:
        print("\n数据预处理成功完成！")
